package collector

import (
	"fmt"

	"dragonfly/common/tcp"
	jww "github.com/spf13/jwalterweatherman"
)

type ConnPool struct {
	count int
	pool  map[string]*tcp.TCPConn
}

func NewConnPool(c *Collector, addrs []string) *ConnPool {
	var cp = &ConnPool{count: 0, pool: make(map[string]*tcp.TCPConn)}

	if err := cp.newConnPool(c, addrs); err != nil {
		jww.ERROR.Println("Connecting failed:", err)
	}
	return cp
}

func (cp *ConnPool) newConnPool(c *Collector, addrs []string) error {
	for _, address := range addrs {
		conn, err := c.Connect(address, nil, nil)
		if err != nil {
			return fmt.Errorf("%v", err)
		}

		if err := cp.Add(address, conn); err != nil {
			jww.ERROR.Println(err)
		}
	}
	return nil
}

func (cp *ConnPool) isExist(addr string) bool {
	_, status := cp.pool[addr]
	return status
}

func (cp *ConnPool) Get(addr string) *tcp.TCPConn {
	if cp.isExist(addr) {
		return cp.pool[addr]
	}
	return nil
}

func (cp *ConnPool) Add(addr string, conn *tcp.TCPConn) error {
	if cp.isExist(addr) {
		return fmt.Errorf("Agent connection is exist. Address: %s", addr)
	}
	cp.pool[addr] = conn
	cp.count += 1

	return nil
}

func (cp *ConnPool) Update(addr string, conn *tcp.TCPConn) error {
	if !cp.isExist(addr) {
		return fmt.Errorf("Agent connection isn't exist. Address: %s", addr)
	}
	cp.pool[addr] = conn
	return nil
}

func (cp *ConnPool) Delete(addr string) error {
	if !cp.isExist(addr) {
		return fmt.Errorf("Agent connection isn't exist. Address: %s", addr)
	}

	cp.Get(addr).Close()
	delete(cp.pool, addr)
	cp.count -= 1
	return nil
}
